package com.xiehua.task.persistence.rocks;

import com.xiehua.task.job.pool.Job;
import com.xiehua.task.persistence.TaskPersistence;
import com.xiehua.task.prop.XiehuaConfig;
import com.xiehua.task.utils.StringUtils;
import org.rocksdb.*;
import org.rocksdb.util.SizeUnit;
import org.tio.utils.json.Json;

import java.io.File;
import java.io.IOException;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static com.xiehua.task.job.pool.Job.JobStatus.fail;
import static com.xiehua.task.job.pool.Job.JobStatus.success;
import static com.xiehua.task.persistence.rocks.RocksFactory.COLUMN_FAMILY_TASK_FAIL;
import static com.xiehua.task.persistence.rocks.RocksFactory.COLUMN_FAMILY_TASK_SUCCESS;
import static com.xiehua.task.prop.XiehuaConfig.PERSISTENCE_PATH;

public class RocksPersitence implements TaskPersistence {

    private Map<String, ColumnFamilyHandle> columnFamilyHandles;

    private RocksDB rocksDB;

    static {
        RocksDB.loadLibrary();
    }

    public RocksPersitence() throws RocksDBException, IOException {
        List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
        //init Options
        Options options = new Options();
        Statistics stats = new Statistics();
        options.setCreateIfMissing(true)
                .setStatistics(stats)
                .setWriteBufferSize(8 * SizeUnit.KB)
                .setMaxWriteBufferNumber(3)
                .setMaxBackgroundCompactions(10)
                .setCompressionType(CompressionType.SNAPPY_COMPRESSION)
                .setCompactionStyle(CompactionStyle.UNIVERSAL);
        Map<String, String> configMap = XiehuaConfig.getInstance();
        String dbPath = RocksPersitence.checkPersitencePath(configMap);
        List<byte[]> cfs = RocksDB.listColumnFamilies(options, dbPath);
        if (cfs.size() > 0) {
            for (byte[] cf : cfs) {
                columnFamilyDescriptors.add(new ColumnFamilyDescriptor(cf, new ColumnFamilyOptions()));
            }
        } else {
            columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, new ColumnFamilyOptions()));
        }

        List<ColumnFamilyHandle> handles = new ArrayList<>();
        DBOptions dbOptions = new DBOptions();
        dbOptions.setCreateIfMissing(true);
        RocksDB rocksDB = RocksDB.open(dbOptions, dbPath, columnFamilyDescriptors, handles);
        Map<String, ColumnFamilyHandle> map = handles.stream().collect(Collectors.toMap(s -> new String(s.getName()), t -> t, (x, y) -> y));
        if (!map.containsKey(COLUMN_FAMILY_TASK_SUCCESS))
            map.put(COLUMN_FAMILY_TASK_SUCCESS, rocksDB.createColumnFamily(new ColumnFamilyDescriptor(COLUMN_FAMILY_TASK_SUCCESS.getBytes(), new ColumnFamilyOptions())));
        if (!map.containsKey(COLUMN_FAMILY_TASK_FAIL))
            map.put(COLUMN_FAMILY_TASK_FAIL, rocksDB.createColumnFamily(new ColumnFamilyDescriptor(COLUMN_FAMILY_TASK_FAIL.getBytes(), new ColumnFamilyOptions())));

        this.columnFamilyHandles = map;
        this.rocksDB = rocksDB;
    }

    @Override
    public void persistence(Job job) throws Exception {
        if (job.getStatus().equals(success)) {
            rocksDB.put(columnFamilyHandles.get(COLUMN_FAMILY_TASK_SUCCESS), job.getJobId().getBytes(), Json.toJson(job).getBytes());
        }
        if (job.getStatus().equals(fail)) {
            rocksDB.put(columnFamilyHandles.get(COLUMN_FAMILY_TASK_FAIL), job.getJobId().getBytes(), Json.toJson(job).getBytes());
        }
    }

    @Override
    public Job queryJob(String taskId) throws Exception {
        byte[] bytes = rocksDB.get(columnFamilyHandles.get(COLUMN_FAMILY_TASK_SUCCESS), taskId.getBytes());
        if (bytes != null && bytes.length > 0) return Json.toBean(new String(bytes), Job.class);
        return null;
    }

    /***
     * 分页查询job
     * @param query
     */
    @Override
    public Job.JobPage queryJob(Job.Query query) {
        return null;
    }

    /**
     * 查询成功的job数量
     *
     * @param start
     * @param end
     **/
    //TODO 待优化
    @Override
    public Integer countSuccessJob(LocalDateTime start, LocalDateTime end, String topic) {
        return countJob(columnFamilyHandles.get(COLUMN_FAMILY_TASK_SUCCESS), start, end, topic);
    }

    /**
     * 查询失败的job数量
     *
     * @param start
     * @param end
     **/
    //TODO 待优化
    @Override
    public Integer countfailJob(LocalDateTime start, LocalDateTime end, String topic) {
        return countJob(columnFamilyHandles.get(COLUMN_FAMILY_TASK_FAIL), start, end, topic);
    }

    /**
     * 统计并过滤job
     **/
    private Integer countJob(ColumnFamilyHandle handle, LocalDateTime start, LocalDateTime end, String topic) {
        RocksIterator iterator = rocksDB.newIterator(handle);
        Integer count = 0;
        for (iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
            Job job = Json.toBean(new String(iterator.value()), Job.class);
            if (topic != null) {
                if (job.getUpdateTime().isAfter(start) && job.getUpdateTime().isBefore(end) && job.getTopic().equals(topic))
                    count++;
            } else {
                if (job.getUpdateTime().isAfter(start) && job.getUpdateTime().isBefore(end)) count++;
            }
        }
        return count;
    }


    public Map<String, ColumnFamilyHandle> getColumnFamilyHandles() {
        return columnFamilyHandles;
    }

    public void setColumnFamilyHandles(Map<String, ColumnFamilyHandle> columnFamilyHandles) {
        this.columnFamilyHandles = columnFamilyHandles;
    }

    public RocksDB getRocksDB() {
        return rocksDB;
    }

    public void setRocksDB(RocksDB rocksDB) {
        this.rocksDB = rocksDB;
    }

    @Override
    public void close() {
        if (rocksDB != null) rocksDB.close();
    }

    /***
     * 路径检查
     * **/
    public static String checkPersitencePath(Map<String, String> configMap) throws IOException {
        if (StringUtils.isEmpty(configMap.get(PERSISTENCE_PATH))) throw new IOException("RocksDB存储路径不能为空");
        LocalDate now = LocalDate.now();
        String persitencePath = configMap.get(PERSISTENCE_PATH) + File.separator + now.getYear() + "_" + now.getMonthValue();
        File persitenceFile = new File(persitencePath);
        if (!persitenceFile.exists()) persitenceFile.mkdir();
        return persitencePath;
    }
}
